import os
import cv2
from tqdm import tqdm # 引入进度条，让处理过程更直观

# --- 1. 配置参数 (高度可自定义) ---

# 输入视频文件所在的文件夹路径
INPUT_VIDEO_FOLDER = '/home/youngqi/Videos' 
# 例如: '/home/youngqi/Videos'，确保这个文件夹存在并包含你的视频文件。

# 抽取图片保存到的目标数据集文件夹路径
# 如果文件夹不存在，脚本会自动创建。
OUTPUT_DATASET_FOLDER = '/home/youngqi/Documents/video_datasheet/images'
# 例如: '/home/youngqi/Documents/datasheet'

# 每秒抽取图片的数量 (Frames Per Second to Extract)
# 例如：
#   - 1：每秒抽取 1 帧
#   - 5：每秒抽取 5 帧
#   - 30：如果视频帧率为 30fps，则每帧都抽取
#   - 0.5：每两秒抽取 1 帧 (注意：结果可能不精确，因为它基于跳过帧数计算)
# 建议：根据你的数据集需求和视频内容调整。对于目标检测，通常 1-5 fps 足够。
FRAMES_PER_SECOND_TO_EXTRACT = 5 

# 接受的视频文件扩展名列表
# 脚本只会处理这些扩展名结尾的文件。
ACCEPTED_VIDEO_EXTENSIONS = ('.mp4', '.avi', '.mov', '.mkv', '.flv', '.m4v')

# 抽取图片的文件名格式
# 格式为 "视频名_序号.jpg"
# 序号将是一个递增的整数，并自动填充前导零以保持文件名长度一致 (例如 00001, 00002)。
# {video_name} 会被替换为视频文件名（不含扩展名）
# {sequence_number:05d} 会被替换为图片序号，并自动填充为5位数（例如 00001）
IMAGE_FILENAME_FORMAT = "{video_name}_{sequence_number:05d}.jpg"

# 是否打印详细的处理信息 (True/False)
VERBOSE = True

# --- 2. 辅助函数 ---

def ensure_directory_exists(path):
    """
    确保指定的目录存在，如果不存在则创建。
    """
    if not os.path.exists(path):
        os.makedirs(path)
        if VERBOSE:
            print(f"已创建目录: {path}")
    else:
        if VERBOSE:
            print(f"目录已存在: {path}")

def process_single_video(video_path, output_dir, frames_per_second, filename_format):
    """
    处理单个视频文件，按指定频率抽取帧并保存。
    """
    video_name_with_ext = os.path.basename(video_path)
    video_name = os.path.splitext(video_name_with_ext)[0]
    
    # 构建当前视频的输出子目录，如果需要按视频名分类保存
    # output_sub_dir = os.path.join(output_dir, video_name)
    # ensure_directory_exists(output_sub_dir) # 如果需要，取消注释此行和下一行
    # current_output_dir = output_sub_dir
    
    # 如果不按视频名分类，所有图片直接保存到 output_dir
    current_output_dir = output_dir

    video = cv2.VideoCapture(video_path)
    if not video.isOpened():
        print(f"错误: 无法打开视频文件 {video_path}")
        return 0

    video_fps = video.get(cv2.CAP_PROP_FPS) # 获取视频的帧率
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) # 获取视频总帧数

    if video_fps == 0:
        print(f"警告: 视频 {video_name_with_ext} 的帧率为 0，可能无法正确抽取图片。跳过。")
        video.release()
        return 0

    # 计算每隔多少帧保存一次图片
    # 如果期望每秒抽取 N 帧，那么 save_step = 视频帧率 / N
    # 例如：视频 30fps，期望每秒 10 帧，则 save_step = 30 / 10 = 3 (每隔 3 帧保存 1 帧)
    # 使用 max(1, ...) 确保 save_step 至少为 1，避免除以零或负数，同时保证至少保存一帧
    if frames_per_second > 0:
        save_step = max(1, int(round(video_fps / frames_per_second)))
    else: # 如果 frames_per_second 为 0 或负数，则默认每帧都保存
        save_step = 1 
        if VERBOSE:
            print(f"警告: 'FRAMES_PER_SECOND_TO_EXTRACT' 设置为 <=0，将抽取视频的所有帧。")


    frame_count = 0
    saved_image_count = 0
    
    if VERBOSE:
        print(f"正在处理视频: {video_name_with_ext} (总帧数: {total_frames}, 帧率: {video_fps:.2f} fps)")
        print(f"每秒抽取帧数: {frames_per_second}, 实际每隔 {save_step} 帧保存一次。")

    # 使用 tqdm 创建进度条
    with tqdm(total=total_frames, unit="frames", desc=f"抽取 {video_name}") as pbar:
        while True:
            ret, frame = video.read()
            if not ret:
                break # 视频读取完毕或出错

            # 检查是否需要保存当前帧
            if frame_count % save_step == 0:
                image_filename = filename_format.format(
                    video_name=video_name,
                    sequence_number=saved_image_count
                )
                output_path = os.path.join(current_output_dir, image_filename)
                
                # 保存图片
                cv2.imwrite(output_path, frame)
                saved_image_count += 1
            
            frame_count += 1
            pbar.update(1) # 更新进度条

    video.release()
    if VERBOSE:
        print(f"视频 {video_name_with_ext} 处理完毕。共保存 {saved_image_count} 张图片。")
    return saved_image_count

# --- 3. 主处理逻辑 ---

def main():
    """
    主函数，扫描视频文件夹，依次处理每个视频。
    """
    ensure_directory_exists(OUTPUT_DATASET_FOLDER)

    total_videos_processed = 0
    total_images_saved = 0

    print(f"\n--- 开始处理视频文件夹: {INPUT_VIDEO_FOLDER} ---")

    for filename in os.listdir(INPUT_VIDEO_FOLDER):
        if filename.lower().endswith(ACCEPTED_VIDEO_EXTENSIONS):
            video_full_path = os.path.join(INPUT_VIDEO_FOLDER, filename)
            
            # 处理单个视频，并累积保存的图片数量
            images_saved_from_video = process_single_video(
                video_full_path, 
                OUTPUT_DATASET_FOLDER, 
                FRAMES_PER_SECOND_TO_EXTRACT, 
                IMAGE_FILENAME_FORMAT
            )
            total_images_saved += images_saved_from_video
            total_videos_processed += 1
        elif VERBOSE:
            print(f"跳过文件: {filename} (不是支持的视频格式)")

    print(f"\n--- 所有视频处理完毕 ---")
    print(f"总计处理视频数量: {total_videos_processed}")
    print(f"总计保存图片数量: {total_images_saved}")
    print(f"所有图片已保存到: {OUTPUT_DATASET_FOLDER}")

# --- 4. 脚本入口 ---
if __name__ == "__main__":
    main()